#!/usr/bin/env python3
"""
WebSocket 客户端示例

演示如何连接到水文预报 WebSocket 服务
"""

import asyncio
import json
import websockets


async def test_websocket():
    """测试 WebSocket 连接"""
    uri = "ws://localhost:8765"
    
    print("="*80)
    print("🔌 连接到 WebSocket 服务器...")
    print("="*80)
    
    async with websockets.connect(uri) as websocket:
        # 接收欢迎消息
        welcome = await websocket.recv()
        print(f"\n📨 收到: {welcome}\n")
        
        # 测试查询列表
        questions = [
            "这次预报的基本信息是什么？",
            "真实预报的洪量、洪峰和峰现时间分别是多少？",
            "对比真实预报和假拟预报的差异",
        ]
        
        for i, question in enumerate(questions, 1):
            print("="*80)
            print(f"测试 {i}: {question}")
            print("="*80)
            
            # 发送问题
            request = {
                "question": question,
                "session_id": "daa0d8a6-dfb3-4887-b798-366f4682c453"
            }
            await websocket.send(json.dumps(request))
            print("✅ 问题已发送")
            
            # 接收响应
            while True:
                response = await websocket.recv()
                data = json.loads(response)
                
                if data["type"] == "processing":
                    print(f"⏳ {data['message']}")
                
                elif data["type"] == "answer":
                    print(f"\n💡 回答:\n{data['answer']}\n")
                    break
                
                elif data["type"] == "error":
                    print(f"\n❌ 错误: {data['error']}\n")
                    break
            
            print("-"*80)
            
            # 等待一下再发送下一个
            if i < len(questions):
                await asyncio.sleep(1)
        
        print("\n" + "="*80)
        print("✅ 测试完成！")
        print("="*80)


if __name__ == "__main__":
    asyncio.run(test_websocket())



